from typing import TypedDict
import uuid

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command

# 定义上下文。这个 State 表示图中每个节点之间的共享状态，当前只有一个字段：some_text，表示一段文本内容。
class State(TypedDict):
    some_text: str

def human_node(state: State):
    print("human node begin")
    #  并将提供的值传递给客户端
    value = interrupt(
        {
            "text_to_revise": state["some_text"]
        }
    )
    return {
        "some_text": value
    }

def say_hello(state: State):
    print("node say_hello output :hello")


# Build the graph
graph_builder = StateGraph(State)
graph_builder.add_node("human_node", human_node)
graph_builder.add_node("say_hello", say_hello)

graph_builder.add_edge(START, "human_node")
graph_builder.add_edge("human_node", "say_hello")

checkpointer = InMemorySaver()

graph = graph_builder.compile(checkpointer=checkpointer)

# Pass a thread ID to the graph to run it.
config = {"configurable": {"thread_id": uuid.uuid4()}}

# Run the graph until the interrupt is hit.
result = graph.invoke({"some_text": "original text"}, config=config)

print(result['__interrupt__'])

# > [
# >    Interrupt(
# >       value={'text_to_revise': 'original text'},
# >       resumable=True,
# >       ns=['human_node:6ce9e64f-edef-fe5d-f7dc-511fa9526960']
# >    )
# > ]

#interrupt 函数将返回通过 Command(resume=value) 提供的恢复值
print(graph.invoke(Command(resume="Edited text"), config=config))
# > {'some_text': 'Edited text'}